"""
Created on Jul 12, 2009

@author: bryan
"""

#===============================================================================
# Imports
#===============================================================================

from scipy import optimize
#from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array, pi
import numpy as np
import JT

#===============================================================================
# Main Class
#===============================================================================
class Optimize():
    """
    This Class will perform 2D-3D image registration, using simulated annealing
    optimization 
    """


    def __init__(self,n_pose=50):
        """
        Constructor
        """
#===============================================================================
#        Data Paths
#===============================================================================
        self._fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        self._inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        self._staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        self._calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
        
        self._roi = [[210,187],[277,370]]
        #self._search_range = 2*array([0.75,0.75,0.75,1.*pi/180,1.25*pi/180,1.25*pi/180])

        # Registration Parameters
        self._n_pose = n_pose
        self._level = 4
        self._method = 'bpc'
        self._debug = True
        self._plot = False
        self._upload_data = False
        
    def Setup(self):
        """
        Setup the DRR and Registration classes for image registration process.
        """
        xraycam = JT.XRayCamera()
        drr = JT.DRR()
        drr.SetXRayCamera(xraycam)
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(self._calFileLat)
        cal.LoadStaFile(self._staFile)
        xraycam.SetExternalCalibration(cal)
        
        self.volume = JT.Volume()
        self.volume.SetVolumeFilename(self._inputVolFileName)
        self.volume.SetOriginalTransform(cal._VolumeT)
        self.volume.SetBodyCenter((3.906, 2.17, -7.378))
        drr.AddVolume(self.volume)
        drr._renWin.SetOffScreenRendering(1)
        
        # Set Fixed image using path
        self.fixedImage = JT.FixedImage()
        self.fixedImage.SetFileName(self._fixedImageName)

        self.reg=JT.Registration()
        self.reg.SetFixedImage(self.fixedImage.GetImage(0))
        self.reg.SetMovingImageGenerator(drr)
        self.reg.SetRegionOfInterest(self._roi)
        
        self.Testing = JT.Testing()
        self.SetMethod('bpc')
        
    
    def Reset(self):
        self.reg.Reset()
    
        
    def FindErr(self,pose):
        """
        Find the magnitude of the error for a given 6 element vector.
        Assumes the the last 3 elements are radians and converts them to degrees
        """
        err = np.sqrt(pose[0]**2 + 
                   pose[1]**2 + 
                   pose[2]**2 +
                   (pose[3]*180/np.pi)**2 +
                   (pose[4]*180/np.pi)**2 + 
                   (pose[5]*180/np.pi)**2)
        return err
    
    def SetMethod(self,method):
        self.reg.SetImageMetric(method)
        
    def anneal(self, search_range):
        """
        """  
        if search_range == None:
            search_range = 2*np.array([0.75,0.75,0.75,1.*pi/180,1.25*pi/180,1.25*pi/180])
            
        #===============================================================================
        #    Setup Optimizer
        #===============================================================================
        pose = self.Testing.GetStartPose(self._n_pose)
        
        # Optimizer specific parameters
        T0=1.0
        Tf=1e-100
        maxiter=20
        maxeval=300
        maxaccept=100
        dwell=30
        feps=1e-100
        boltzmann=5e-4
        learn_rate=0.5  #Not used in Fast schedule
        quench=1.01
        m=1.3
        n=.7
        schedule='fast'
        
        # Set upper and lower bounds
        lowerBounds = pose - search_range/2
        upperBounds = pose + search_range/2
        
        x0 = pose
        
        # Set cost function
        f=self.reg.GetValue
        # Additional Arguments for the cost function
        debug = True
        plot = False

        (xmin, Jmin, T, feval, iters, accept, retval) = optimize.anneal(
                                    f, x0, args=(debug,plot),
                                    full_output=True, 
                                    T0=T0, Tf=Tf, lower=lowerBounds, 
                                    upper=upperBounds, m=m, n=n,
                                    maxiter=maxiter, maxeval=maxeval, 
                                    maxaccept=maxaccept, dwell=dwell, 
                                    feps=feps, schedule=schedule,
                                    boltzmann=boltzmann, quench=quench, 
                                    learn_rate=learn_rate)
        
        print "Solution: ", xmin
        err = self.FindErr(xmin)

        return err
    
    def fmin(self,search_range=None):
        """
        fmin search
        """
        
        if search_range == None:
            search_range = 2*np.array([0.5,0.5,0.5,1.*pi/180,1.*pi/180,1.*pi/180])
          
        # Set the cost function
        f=self.reg.GetValue
        
        # Additional Arguments for the cost function
        debug = True
        plot = False
        
        # Set the upper and lower bounds based on the initial guess
        pose = self.Testing.GetStartPose(self._n_pose)
        lowerBounds = pose - search_range/2
        upperBounds = pose + search_range/2
        
        (xopt, fopt, iter, funcalls, warnflag) = optimize.fmin(
                                    f, x0 = pose, xtol = 0.01, maxfun = 300, 
                                    maxiter = 200, full_output = 1, disp = 1)
        err = self.FindErr(xopt)
        
        
    def brent(self):
        """
        brent bracketing search
        """
        # Set the cost function
        def myFunc(x):
            pose = [.1,.1,.1,x,.1,.1]
            return self.reg.GetValue(pose)
        
        f = myFunc
        # Additional Arguments for the cost function
        debug = True
        plot = False
        
        # Set the upper and lower bounds based on the initial guess
        a = -.1
        b = .1
        c = .25
        brack = (a,c)
        
        (xopt, fopt, iter, funcalls) = optimize.brent(
                                    f, brack = brack, tol=1.48e-4,
                                    maxiter = 200, full_output = 1)
        
        print "xopt: ", xopt, ", fopt: ", fopt
        return xopt
    
    def fmin_bfgs(self):
        """
        """
         # Set the cost function
        f=self.reg.GetValue
        
        # Set the upper and lower bounds based on the initial guess
        pose = self.Testing.GetStartPose(self._n_pose)
        pose = [4,4,4,0,0,0]
        
        (xopt,fopt, gopt, Hopt, func_calls, grad_calls, warnflag,allvecs) = \
        optimize.fmin_bfgs(f, pose, fprime=None, args=(), gtol=1e-5,
                  epsilon=1e-4, maxiter=200, full_output=1, disp=1,
                  retall=1, callback=None)
    
    def fmin_powell(self):
        """
        """
        # Set the cost function
        f=self.reg.GetValue
        
        # Set the upper and lower bounds based on the initial guess
        pose = self.Testing.GetStartPose(self._n_pose)
        #pose = [4,4,4,0,0,0]
        
        direc = 1*np.eye(6)
        direc[3:6,3:6] = 0.02*direc[3:6,3:6]
        
        (xopt, fopt, xi, direc, iter, funcalls, allvecs) = \
        optimize.fmin_powell(f, pose, args=(), xtol=1e-1, ftol=1e-1, maxiter=100,
                maxfun=100, full_output=1, disp=1, retall=1, callback=None,
                direc=direc)
        
    def grad_des(self):
        """
        Gradient Descent optimizer (currently located in JT module)
        """
        # Set the cost function
        f=self.reg
        
        # Set the starting pose
        pose = self.Testing.GetStartPose(self._n_pose)
        
        gd = JT.OptimizerGD(f,startx=pose,abstol=1e-3,verbose=1,stepf=1.2,
                            startstep=2,minstep=1e-3)
        
        while not gd.hasConverged():
            gd.makeStep()
            
        err = gd.getX()
        print err
        return err
       
    def ITKRegStepGradDes(self):
        """
        ITK implementation of the regular step gradient descent optimizer.
        Translated into python from C++.
        """
        from JT.ITKOptimizers import RegularStepGradientDescentOptimizer
        gd = RegularStepGradientDescentOptimizer()
        
        # Set the starting pose
        pose = self.Testing.GetStartPose(self._n_pose)
        gd.SetInitialPosition(pose)
        
        # Set Scales 1/1000 for translations, 1 for rotations
        scales = [1.0,1.0,1.0,500.0,500.0,500.0]
        gd.SetScales(scales)
        
        # Set the cost function
        gd.SetCostFunction(self.reg)
        
        # Modify Optimizer Parameters
        gd.SetMaximumStepLength(1)
        
        
        # Start the optimizer
        gd._debug = True
        gd.StartOptimization()
        
        # Print Final State    
        gd.PrintSelf()
        
        return self.FindErr( gd.GetCurrentPosition() )

if __name__ == '__main__':
    Opt = Optimize(n_pose=60)
    Opt.Setup()
    err = Opt.ITKRegStepGradDes() #optionally pass the search range as an argument
    #Opt.Reset()
    
    
        